package com.wangshb.agent.demo.runtime;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashMap;


/**
 *
 * 统计运行时间
 *
 * @author wangsb12
 */
public class MyTransformer implements ClassFileTransformer {

    final static String prefix = "\nlong startTime = System.currentTimeMillis();\n";
    final static String postfix = "\nlong endTime = System.currentTimeMillis();\n";

    /** 拦截的包路径 */
    final  static String INVOKE_PACKAGE = "com.wangshb";

    /** 忽略某些方法，不需要统计运行时间 */
    final static HashMap<String, Object> ignoreMethodMap = new HashMap<String, Object>(){
        {
            this.put("equals", new Object());
            this.put("finalize", new Object());
            this.put("toString", new Object());
            this.put("getClass", new Object());
            this.put("notifyAll", new Object());
            this.put("hashCode", new Object());
            this.put("wait", new Object());
            this.put("notify", new Object());
            this.put("main", new Object());
            this.put("clone", new Object());
        }
    };


    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        className = className.replace("/", ".");

        if (!className.contains(INVOKE_PACKAGE)) {
            return classfileBuffer;
        }
        CtClass ctclass ;
        try {
            ClassPool pool = ClassPool.getDefault();

            // 将加载器加入到classPath
            pool.insertClassPath(new LoaderClassPath(loader));

            // 使用全称,用于取得字节码类<使用javassist>
            ctclass = pool.get(className);
            CtMethod[] methods = ctclass.getMethods();
            for (CtMethod ctmethod : methods) {
                String methodName = ctmethod.getName();
                if (ignoreMethodMap.containsKey(methodName)) {
                    continue;
                }
                String outputStr = "\nSystem.out.println(\"this method " + methodName
                        + " cost:\" +(endTime - startTime) +\"ms.\");";

                // 新定义一个方法叫做比如sayHello$old, 同时修改原方法名称
                String newMethodName = methodName + "$old";
                ctmethod.setName(newMethodName);
                // 创建新的方法，复制原来的方法，名字为原来的名字
                CtMethod newMethod = CtNewMethod.copy(ctmethod, methodName, ctclass, null);

                // 构建新的方法体
                StringBuilder bodyStr = new StringBuilder();
                bodyStr.append("{");
                bodyStr.append(prefix);
                // 调用原有代码，类似于method();($$)表示所有的参数
                bodyStr.append(newMethodName + "($$);\n");
                bodyStr.append(postfix);
                bodyStr.append(outputStr);
                bodyStr.append("}");
                // 替换新方法
                newMethod.setBody(bodyStr.toString());

                // 增加新方法
                ctclass.addMethod(newMethod);
            }
            return ctclass.toBytecode();
        } catch (NotFoundException e) {
            e.printStackTrace();
        } catch (CannotCompileException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return classfileBuffer;
    }
}
